#encoding: utf-8


import cv2
import numpy as np

import torch
from torch import nn
import torchvision.utils as vutils
import torchvision.transforms as T
from torch.nn import functional as F
import torch.nn.utils.weight_norm as weightNorm
from torch.nn.parameter import Parameter

import sys
import copy
sys.path.append('..')

from .backbone.resnet import ResNet, BasicBlock, Bottleneck
from .backbone.PDecouple_swin import PSwinTransformer
from .backbone.PDecouple_resnet import Presnet

BN = None

def weights_init_kaiming(m):
    classname = m.__class__.__name__
    if classname.find('Linear') != -1:
        nn.init.kaiming_normal_(m.weight, a=0, mode='fan_out')
        if m.bias is not None:
          nn.init.constant_(m.bias, 0.0)
    elif classname.find('Conv') != -1:
        nn.init.kaiming_normal_(m.weight, a=0, mode='fan_in')
        if m.bias is not None:
            nn.init.constant_(m.bias, 0.0)
    elif classname.find('BatchNorm') != -1:
        if m.affine:
            nn.init.constant_(m.weight, 1.0)
            nn.init.constant_(m.bias, 0.0)

def weights_init_reduction(m):
    classname = m.__class__.__name__
    if classname.find('Conv') != -1:
        nn.init.constant_(m.weight, 1.0 / m.weight.size(0))
        nn.init.constant_(m.bias, 0.0)
    elif classname.find('BatchNorm') != -1:
        if m.affine:
            nn.init.constant_(m.weight, 1.0)
            nn.init.constant_(m.bias, 0.0)

def weights_init_classifier(m):
    classname = m.__class__.__name__
    if classname.find('Linear') != -1:
        nn.init.normal_(m.weight, std=0.001)
        if m.bias:
            nn.init.constant_(m.bias, 0.0)


class Baseline(nn.Module):
    in_planes = 2048
    reduction = 1
    def __init__(self, cfg,num_classes, last_stride, model_path, neck, neck_feat, model_name, pretrain_choice):
        super(Baseline, self).__init__()
        self.tdeep=cfg.MODEL.TDeep

        global BN

        if model_name == 'resnet18':
            self.in_planes = 512
            self.reduction = 2
            self.base = ResNet(last_stride=last_stride,
                               block=BasicBlock,
                               layers=[2, 2, 2, 2], bn_group=bn_group, bn_var_mode=bn_var_mode)
        elif model_name == 'resnet34':
            self.in_planes = 512
            self.base = ResNet(last_stride=last_stride,
                               block=BasicBlock,
                               layers=[3, 4, 6, 3], bn_group=bn_group, bn_var_mode=bn_var_mode)
        elif model_name == 'resnet50':
            self.base = ResNet(last_stride=last_stride,
                               block=Bottleneck,
                               layers=[3, 4, 6, 3], bn_group=bn_group, bn_var_mode=bn_var_mode)
        elif model_name == 'resnet101':
            self.base = ResNet(last_stride=last_stride,
                               block=Bottleneck,
                               layers=[3, 4, 23, 3], bn_group=bn_group, bn_var_mode=bn_var_mode)
        elif model_name == 'resnet152':
            self.base = ResNet(last_stride=last_stride,
                               block=Bottleneck,
                               layers=[3, 8, 36, 3], bn_group=bn_group, bn_var_mode=bn_var_mode)

        elif model_name == 'Pres50':
            self.base = Presnet(model_path=model_path,last_stride=last_stride,pre_trained=pretrain_choice,)
            self.in_planes = 2048
            self.reduction = cfg.MODEL.nssd_feature_dim
   
            # self.in_planes2 = int(2048)
            # self.reduction2 = int(cfg.MODEL.nssd_feature_dim)        
        self.model_name=model_name

        if pretrain_choice == 'imagenet':
            if model_name == 'PSwin' or model_name=="Pres50":
                pass
            else:
                self.base.load_param(model_path)
                print('Loading pretrained ImageNet model {}......'.format(model_path))

        self.update_param()
        self.gap = nn.AdaptiveAvgPool2d(1)
        self.aap = nn.AdaptiveAvgPool2d(2)

        self.num_classes = num_classes
        self.neck = neck
        self.neck_feat = neck_feat
        self.cam=cfg.MODEL.nssd_grad_cam

        if self.neck == 'bnneck':
            if self.model_name=="PSwin" or self.model_name == 'Pres50':
                pass
            else:
                self.bnneck = nn.BatchNorm1d(int(self.in_planes/self.reduction))
                self.bnneck.bias.requires_grad_(False)  # no shift
                self.bnneck.apply(weights_init_kaiming)

                self.classifier = nn.Linear(int(self.in_planes/self.reduction), self.num_classes, bias=False)
                self.classifier.apply(weights_init_classifier)

        if self.model_name == 'Pres50':
            self.classifier = nn.Linear(int(self.reduction), self.num_classes, bias=False)
            self.classifier.apply(weights_init_classifier)
            self.reduction_layer1 = nn.Linear(int(self.in_planes), int(self.reduction), bias=False)
            self.reduction_layer1.apply(weights_init_classifier)
            self.bnneck1 = nn.BatchNorm1d(int(self.reduction))
            self.bnneck1.bias.requires_grad_(False)
            self.bnneck1.apply(weights_init_kaiming)
            self.bn_layers = nn.ModuleList()
            self.reduction_layers = nn.ModuleList()
            self.cls_layers = nn.ModuleList()
            self.bn_layers.append(self.bnneck1)
            self.reduction_layers.append(self.reduction_layer1)
            self.cls_layers.append(self.classifier)
 


    def forward(self, x, epoch=0,writer=None, n_iter=None):

        if self.model_name == 'Pres50':
            base_feats= self.base(x,self.tdeep)
            global_feat = self.reduction_layers[0](base_feats[0])
            if self.model_name == 'Pres50' and self.tdeep==0:
                all_feats=[global_feat]
       

        if self.model_name == 'Pres50':
            if self.neck == 'no':
                feat = all_feats
            elif self.neck == 'bnneck':
                bn_all_feats=[]
                if self.model_name == 'Pres50' and self.tdeep==0:
                    tmp_feat = self.bn_layers[0](all_feats[0])
                    bn_all_feats.append(tmp_feat)
              

        if self.training:
            if self.model_name == 'Pres50':
                all_scores=[]              
                if self.model_name == 'Pres50' and self.tdeep==0:
                    tmp_score = self.cls_layers[0](bn_all_feats[0])
                    all_scores.append(tmp_score)
     
                return all_scores,all_feats
            else:
                cls_score = self.classifier(feat)
                return cls_score, global_feat  # global feature for triplet loss
        else:
            # print("evaluate here")
            if self.neck_feat == 'after':
                # return feat
                if self.model_name == 'Pres50':
                    if self.model_name == 'Pres50' and self.tdeep==0:
                        return bn_all_feats[0] #, bn_all_feats[0],bn_all_feats[0]
                    
            else:
                if self.model_name == 'Pres50':
                    if self.model_name == 'Pres50' and self.tdeep==0:
                        return all_feats[0] # , all_feats[0],all_feats[0]

                else:
                    return global_feat

    def load_param(self, trained_path):
        param_dict = torch.load(trained_path)['state_dict']
        for key in param_dict:
            key_cut = key[7:]
            if 'classifier' in key_cut:
                continue
            self.state_dict()[key_cut].copy_(param_dict[key])

    def load_param_swin(self, trained_path):
        param_dict=torch.load(trained_path, 'cpu')['model']
        for key in param_dict:
            new_key = "base."+key
            if 'head' in key:
                continue
            try:
              self.state_dict()[new_key].copy_(param_dict[key])
            except:
              print(key)
              exit()

    def update_param(self):
        pass




